package uk.ac.cam.cl.oscar.sift;

import java.util.Vector;
import java.awt.image.BufferedImage;

public class Octave
{
    // Our two stacks of images - the blurred images, and the
    // difference-of-gaussians
    private Vector<BufferedImage> blurredImages, doGImages;
    // The gradients of the image need to be calculated later on
    // so we'll store stuff about them here
    private Vector<BufferedImage> lengths, directions;

    public Octave()
    {
        blurredImages = new Vector<BufferedImage>();
        doGImages = new Vector<BufferedImage>();
    }

    /**
     * Adds an array of BufferedImages to the stack of blurred images.
     */
    public void addBlurredImages(BufferedImage[] images)
    {
        for (int n=0; n<images.length; n++) {
            blurredImages.addElement(images[n]);
        }
    }

	/** Empties the blurred images array */
	public void clearBlurredImages () {
		blurredImages.clear();
	}

    /**
     * Adds an array of BufferedImages to the stack of DoG images.
     */
    public void addDoGImages(BufferedImage[] images)
    {
        for (int n=0; n<images.length; n++) {
            doGImages.addElement(images[n]);
        }
    }

	public int doGLevels () {
		return doGImages.size();
	}

    /**
     * Returns the blurred image at index n on the stack.
     */
    public BufferedImage getBlurredImage(int n)
    {
        return blurredImages.elementAt(n);
    }

    public int getBlurredImageCount() { return blurredImages.size(); }

    public Vector<Peak> findPeaks(double threshold)
    {
        Vector<Peak> peaks = new Vector<Peak>();

        BufferedImage current, above, below;
        
        int actualThreshold = (int)Math.round(256 * threshold);

        // For each image in the stack of DoG images
        for (int level=1; level<(doGImages.size()-1); level++) {
            current = doGImages.elementAt(level);
            below = doGImages.elementAt(level-1);
            above = doGImages.elementAt(level+1);

            boolean ismax = true;
            boolean ismin = true;

            // For every pixel in the image (excluding border pixels)
            // try and see if it's a local minimum or local maximum
            // Only the blue (lowest byte) channel carries useful data, so we mask that and use it straight away 
            
            for (int y=1; y<(current.getHeight()-1); y++) 
            {
                for (int x=1; x<(current.getWidth()-1); x++) 
                {
                   int mag = current.getRGB(x,y) & 0xFF;

                   if (mag > actualThreshold) {
                       /* it's a deep loop. so kill me. */
                       for (int dx = -1; dx <= 1 && ismax; dx++) {
                     		for (int dy = -1; dy <= 1 && ismax; dy++) {
                     			if (dx != 0 || dy != 0) {
                     				if ( (current.getRGB(x + dx, y + dy)&0xFF) > mag) {
                     					ismax = false;
                     					break;
                     				}
                     				else if ( (above.getRGB(x + dx, y + dy)&0xFF) > mag) {
                     					ismax = false;
                     					break;
                     				}
                     				else if ( (below.getRGB(x + dx, y + dy)&0xFF) > mag) {
                     					ismax = false;
                     					break;
                     				}
                     			}
                     		}
                       }
                       
                       if (ismax) {
                          peaks.add(new Peak(x,y,level));
                       } else {
                           
							for (int dx = -1; dx <= 1 && ismin; dx++) {
								for (int dy = -1; dy <= 1 && ismin; dy++) {
									if (dx != 0 || dy != 0) {
										if ( (current.getRGB(x + dx, y + dy)&0xFF) < mag) {
                     					ismin = false;
                     					break;
                     				}
                     				else if ( (above.getRGB(x + dx, y + dy)&0xFF) < mag) {
                     					ismin = false;
                     					break;
                     				}
                     				else if ( (below.getRGB(x + dx, y + dy)&0xFF) < mag) {
                     					ismin = false;
                     					break;
                     				}
									}
								}
						   }
                           
                           if (ismin) {
                               peaks.add(new Peak(x,y,level));
                           }
                       }
                       ismax = true;
                       ismin = true;
                   }
                }
            }
        }
        return peaks;
    }

    public Vector<Peak> filterPeaks(Vector<Peak> peaks, double edgeRatio)
    {
        Vector<Peak> filtered = new Vector<Peak>();

        for (Peak pk : peaks) {
            // First off, we have to eliminate edge responses.  According to (Lowe, 2004),
            // we can do this by calculating the elements of the Hessian matrix computed
            // at the location and scale of the keypoint
            int x = pk.getX(), y = pk.getY();
            double d_XX, d_YY, d_XY; // Has nothing at all to do with chromosomes
            BufferedImage curlevel = doGImages.elementAt(pk.getLevel());

            // Calculating the Hessian elements
            d_XX = curlevel.getRGB(x+1,y) + curlevel.getRGB(x-1,y) - 2.0 * curlevel.getRGB(x,y);
            d_YY = curlevel.getRGB(x,y+1) + curlevel.getRGB(x,y-1) - 2.0 * curlevel.getRGB(x,y);
            d_XY = 0.25 * ((curlevel.getRGB(x+1,y+1) - curlevel.getRGB(x+1,y-1)) -
                   (curlevel.getRGB(x-1,y+1) - curlevel.getRGB(x-1,y-1)));

            double traceSq = d_XX + d_YY; // The trace of the hessian matrix...
            traceSq *= traceSq; // Squared!
            double det = d_XX * d_YY - (d_XY * d_XY); // Determinant of the hessian matrix

            // "Therefore, to check the ratio of principal curvatures is below some threshold,
            // r [edgeRatio] we only need to check Tr(H)^2 / Det(H) < (r+1)^2 / r" (Lowe, 2004)
            if ((traceSq / det) >= ((edgeRatio + 1.0) / edgeRatio)) {
                // Ok, so it isn't an edge.  Add it to our filtered list
                filtered.addElement(pk);
            }
        }
        return filtered;
    }

    public void makeDescriptors()
    {
    	
        lengths = new Vector<BufferedImage>();
        directions = new Vector<BufferedImage>();

        for (int level=1; level<(doGImages.size()-1); level++) {
            BufferedImage current = blurredImages.elementAt(level);
            int height = current.getHeight();
            int width = current.getWidth();
            BufferedImage length = new BufferedImage(width, height,
                                                     BufferedImage.TYPE_INT_ARGB);
            BufferedImage direction = new BufferedImage(width,
                                                        height,
                                                        BufferedImage.TYPE_INT_ARGB);

            for (int y=1; y<(height-1); y++) {
                for (int x=1; x<(width-1); x++) {
                    // Length of a vector is sqrt(x^2 + y^2)
                    // Corrected 22/02/08 - Could have overflew.
                    double rgb = Math.sqrt(Math.pow(current.getRGB(x+1,y) - current.getRGB(x-1,y), 2.0) +
                                          Math.pow(current.getRGB(x, y+1) - current.getRGB(x,y-1), 2.0));
                    rgb = (rgb > 16777215.0) ? 1677215.0 : rgb;

                    length.setRGB(x,y,(int)rgb);

                    // Direction between -PI and PI, normalised for RGB values
                    // in the range 0 to 16777215
                    rgb = Math.atan2(current.getRGB(x,y+1) - current.getRGB(x, y-1),
                                         current.getRGB(x+1, y) - current.getRGB(x-1, y));
                    rgb /= Math.PI;
                    rgb *= 8388607;
                    rgb += 8388607;
                    direction.setRGB(x,y,(int)rgb);
                }
            }
            
            lengths.addElement(length);
            directions.addElement(direction);
        }
    }

    public Vector<Keypoint> generateKeypoints(Vector<Peak> peaks, double imgScale, double octaveSigma, int scaleCount)
    {
        Vector<Keypoint> keypoints = new Vector<Keypoint>();

        for (Peak pk : peaks) {
            Vector<Keypoint> tmp = new Vector<Keypoint>();

            double kpScale = octaveSigma * Math.pow(2.0, pk.getLevel() / scaleCount);

            double sigma = 3.0 *kpScale;
            int radius = (int)(3.0 * sigma / 2.0 + 0.5);
            int radiusSq = radius * radius;

            BufferedImage length = lengths.elementAt(pk.getLevel() - 1); 
            /* - 1 since the length corresponding to level 1 gets added to an
               empty stack, and therefore will be at point 0 */
            			
            BufferedImage direction = directions.elementAt(pk.getLevel() - 1);

            int xMin = Math.max(pk.getX() - radius, 1);
            int xMax = Math.min(pk.getX() + radius, length.getWidth()-1);
            int yMin = Math.max(pk.getY() - radius, 1);
            int yMax = Math.min(pk.getY() + radius, length.getHeight()-1);

            double gaussianSigmaFactor = 2.0 * sigma * sigma;

            double[] bins = new double[36];

            // Here we build a direction histogram with 36 bins, each representing a
            // 10 degree segment
            for (int y=yMin; y<yMax; y++) {
                for (int x=xMin; x<xMax; x++) {
                    int relX = x - pk.getX();
                    int relY = y - pk.getY();

                    // If it lies within the circle
                    if ((relX * relX) + (relY * relY) <= radiusSq) {
                        // Gaussian weight factor = e^{- r^2 / 2*sigma^2}
                        double gaussianWeight = Math.pow(Math.E, -((relX * relX) + (relY * relY) / gaussianSigmaFactor));

                        double angle = ((direction.getRGB(x,y)-8388607)/8388607 + Math.PI) / (2.0 * Math.PI);
                        angle *= 36;
                        int bin = (int)angle;
                        if (bin == 36) bin = 0;

                        bins[bin] += length.getRGB(x,y) * gaussianWeight;
                    }
                }
            }

            // Find the highest bar on the histogram
            double maxGrad = bins[0];
            int maxBin = 0;
            for (int b=1; b<36; b++) {
                if (bins[b] > maxGrad) {
                    maxGrad = bins[b];
                    maxBin = b;
                }
            }

            // Estimate the peak value and correction value
            double maxPeakValue, left, right;
            if (maxBin == 0) {
                left = bins[35];
            } else {
                left = bins[maxBin-1];
            }
            if (maxBin == 35) {
                right = bins[0];
            } else {
                right = bins[maxBin+1];
            }
            double a = ((left + right) - 2.0 * bins[maxBin]) / 2.0;
            double c = (((left - bins[maxBin]) / a) - 1.0) / 2.0;
            double b = bins[maxBin] - c * c * a;
            maxPeakValue = b;

            // We have the peak value, so we need to find other orientations
            boolean[] binIsKeypoint = new boolean[36];
            for (int i=0; i<36; i++) {
                binIsKeypoint[i] = false;

                // Max Peak is obviously a peak value
                if (i == maxBin) {
                    binIsKeypoint[i] = true;
                } else if (bins[i] >= (0.8 * maxPeakValue)) {
                   int leftI = (i==0) ? 35 : i-1;
                   int rightI = (i==35) ? 0 : i+1;
                   if (bins[i] > bins[leftI] && bins[i] > bins[rightI]) {
                       binIsKeypoint[i] = true; // It's a local peak
                   }
                }
            }

            double oneBinRad = (2.0 * Math.PI) / 36;

            for (int i=0; i<36; i++) {
                  if (binIsKeypoint[i]) {
                      int leftI = (i==0) ? 35 : i-1;
                      int rightI = (i==35) ? 0 : i+1;
                      left = bins[leftI];
                      right = bins[rightI];

                      double degreeCorrection;
                      a = ((left + right) - 2.0 * bins[i]) / 2.0;
                      c = (((left - bins[i]) / a) - 1.0) / 2.0;
                      b = bins[i] - c * c * a;
                      degreeCorrection = c;

                      double degree = (i + degreeCorrection) * oneBinRad - Math.PI;

                      if (degree < -Math.PI) {
                          degree += 2.0 * Math.PI;
                      } else if (degree > Math.PI) {
                          degree -= 2.0 * Math.PI;
                      }

                      Keypoint kp = new Keypoint(blurredImages.elementAt(pk.getLevel()), pk.getX(), pk.getY(), imgScale, kpScale, degree);
                      tmp.addElement(kp);
                  }
            }
            tmp = createDescriptors(tmp, lengths.elementAt(pk.getLevel()-1), directions.elementAt(pk.getLevel()-1), 2.0, 4, 8, 0.2);

            for (Keypoint kp : tmp) {
                kp.setX(kp.getX() * kp.getImgScale());
                kp.setY(kp.getY() * kp.getImgScale());
                kp.setScale(kp.getScale() * kp.getImgScale());

                keypoints.addElement(kp);
            }
        }
        return keypoints;
    }

    public Vector<Keypoint> createDescriptors(Vector<Keypoint> kplist, BufferedImage length, BufferedImage direction, double scale, int dim, int count, double cap)
    {
        double scalefactor = scale * kplist.elementAt(0).getScale();
        int radius = (int) (((dim + 1) / 2) * Math.sqrt(2.0) * scalefactor + 0.5);

        Vector<Keypoint> result = new Vector<Keypoint>();
 
        for (Keypoint kp : kplist) {
            KeypointDescriptor kpd = new KeypointDescriptor(dim, dim, count);

            for (int y=-radius; y<radius; y++) {
                for (int x=-radius; x<radius; x++) {
                    // Rotate y and x
                    double newy = Math.sin(kp.getOrientation()) * x + Math.cos(kp.getOrientation()) * y;
                    double newx = Math.cos(kp.getOrientation()) * x - Math.sin(kp.getOrientation()) * y;

                    // Scale the rotated x and y
                    newy /= scalefactor;
                    newx /= scalefactor;

                    // Consider all (newx, newy) that are in the square defined by the corners
                    // (-dim/2 - 0.5, -dim/2-0.5) and (dim/2 + 0.5, dim/2 + 0.5)
                    if ((y < ((dim / 2.0) + 0.5)) && (y > (-(dim / 2.0) - 0.5)) &&
                        (x < ((dim / 2.0) + 0.5)) && (x > (-(dim / 2.0) - 0.5))) {
                        int currentx = (int)(x + kp.getX() + 0.5);
                        int currenty = (int)(y + kp.getY() + 0.5);

						/* Why where these lines inside the next statement? - DMack */
						currentx += (dim / 2.0) - 0.5;
                        currenty += (dim / 2.0) - 0.5;

                        // If the location isn't in the bounds of our precalculated descriptor maps
                        // then we don't give a damn
                        if ((currentx > 1) && (currentx < (length.getWidth()-1)) &&
                            (currenty > 1) && (currenty < (length.getHeight()-1))) {
                            

                            // Build linear interpolation weights:
                            // A B
                            // C D
                            //
                            // The keypoint is located between A, B, C and D.
                            int[] xIdx = new int[2];
                            int[] yIdx = new int[2];
                            int[] dirIdx = new int[2];
                            double[] xWeight = new double[2];
                            double[] yWeight = new double[2];
                            double[] dirWeight = new double[2];

                            if (newx >= 0) {
                                xIdx[0] = (int) newx;
                                xWeight[0] = (1.0 - (newx - xIdx[0]));
                            }
                            if (newy >= 0) {
                                yIdx[0] = (int) newy;
                                yWeight[0] = (1.0 - (newy - yIdx[0]));
                            }
                            if (newx < (dim - 1)) {
                                xIdx[1] = (int) (newx + 1.0);
                                xWeight[1] = newx - xIdx[1] + 1.0;
                            }
                            if (newy < (dim - 1)) {
                                yIdx[1] = (int) (newy + 1.0);
                                yWeight[1] = newy - yIdx[1] + 1.0;
                            }
                            
                            double dir = (direction.getRGB(currentx,currenty) - 8388607) / 8388607 - kp.getOrientation();
                            if (dir <= -Math.PI) {
                                dir += Math.PI;
                            }
                            if (dir > Math.PI) {
                                dir -= Math.PI;
                            }
                            double idxDir = (dir * count) / (2.0 * Math.PI);
                            if (idxDir < 0.0) {
                                idxDir += count;
                            }
                            dirIdx[0] = (int)idxDir;
                            dirIdx[1] = (dirIdx[0] + 1) % count;
                            dirWeight[0] = 1.0 - idxDir - dirIdx[0];
                            dirWeight[1] = idxDir - dirIdx[0];

                            double sigma2Sq = 2.0 * (dim / 2.0) * (dim / 2.0);
                            double mag = Math.pow (Math.E, -(newx * newx + newy * newy) / sigma2Sq) * length.getRGB(currentx, currenty);

                            for (int iy=0; iy<2; iy++) {
                                for (int ix=0; ix<2; ix++) {
                                    for (int id=0; id<2; id++) {
                                        kpd.set(xIdx[ix], yIdx[iy], dirIdx[id], 
                                                         kpd.get(xIdx[ix], yIdx[iy], dirIdx[id]) +
                                                         xWeight[ix] * yWeight[iy] * dirWeight[id] * mag);
                                    }
                                }
                            }
                            kp.addDescriptor(kpd);
                        }
                    }
                }
            }
            KeypointDescriptor kpdn = kp.getDescriptor();
            if (kpdn != null) {
                // Normalise the descriptors.
                double norm = 0.0;
                for (int n=0; n<kpdn.getLinearDim(); n++) {
                    norm += Math.pow(kpdn.linearGet(n), 2.0);
                }
                norm = Math.sqrt(norm);

                for (int n=0; n<kpdn.getLinearDim(); n++) {
                    kpdn.linearSet(n, kpdn.linearGet(n) / norm);

                    if (kpdn.linearGet(n) > cap) {
                        kpdn.linearSet(n, cap);
                    }
                }

                norm = 0.0;
                for (int n=0; n<kpdn.getLinearDim(); n++) {
                    norm += Math.pow(kpdn.linearGet(n), 2.0);
                }
                norm = Math.sqrt(norm);

                for (int n=0; n<kpdn.getLinearDim(); n++) {
                    kpdn.linearSet(n, kpdn.linearGet(n) / norm);
                }
                kp.addDescriptor(kpdn);

                result.add(kp);
            }
        }
        return result;
    }
}
